#include "srpq_controller/AdvanceControllers/BalanceController/BalanceController.hpp"
#include <common/Utilities/Utilities_print.h>
#include <common/Utilities/Timer.h>

template <typename T>
BalanceController<T>::BalanceController(){
    _body_pos_task = new BodyPosTask<T>();
    _body_ori_task = new BodyOriTask<T>();
    _foot_task = new FootTask<T>();
    _contact_task = new ContactTask<T>();

    // task arrange
    task_list.push_back(_foot_task);
    task_list.push_back(_body_pos_task);
    task_list.push_back(_body_ori_task);
}

template <typename T>
BalanceController<T>::~BalanceController(){}

template <typename T>
void BalanceController<T>::run(void* input, ControlFSMData<T>& data) {
    ++_iter;
    // Compute control output
    _ComputeBC(input, data);

    // Update Leg Command
    _UpdateLegCMD(data);
}


template <typename T>
void BalanceController<T>::_ComputeBC(void* input, ControlFSMData<T>& data){
    // Get desire state
    _input_data = static_cast<BalanceCtrlData<T>* >(input);

    // Get current state
    auto& seResult = data._stateEstimator->getResult();

    Vec3<T> qDes, qdDes;
    for(int leg = 0; leg < 4; leg ++){
        // contact state
        if(seResult.contactEstimate(leg) > 0.){
            _contact_task->run(data, _input_data, leg);

            // output
            // _pid_data.qDes[leg] = _contact_task->rs.qDes;
            _pid_data.qdDes[leg] = _contact_task->rs.qdDes;
            // _pid_data.kpJoint[leg] << 10,0,0,0,10,0,0,0,10;
            _pid_data.kdJoint[leg] << 0.2,0,0,0,0.2,0,0,0,0.2;           

        }

        // swing state
        else{
            _foot_task->run(data, _input_data, leg);
            _body_pos_task->run(data, _input_data, leg);
            _body_ori_task->run(data, _input_data, leg);

            // weight
            T w = 0.6;
            qDes = Vec3<T>::Zero();
            qdDes = Vec3<T>::Zero();
            for(int i = 0; i < task_list.size(); i ++){
                if(!i){
                    qDes = task_list[i]->rs.qDes;
                    qdDes = task_list[i]->rs.qDes;
                }
                else{
                    qDes = qDes * w + task_list[i]->rs.qDes * (1-w);
                    qdDes = qdDes * w + task_list[i]->rs.qdDes * (1-w);
                }
            }

            // output
            _pid_data.qDes[leg] = qDes;
            _pid_data.qdDes[leg] = qdDes;
            _pid_data.kpJoint[leg] << 12,0,0,0,12,0,0,0,12;
            _pid_data.kdJoint[leg] << 0.5,0,0,0,0.5,0,0,0,0.5;

        }

    }
}

template <typename T>
void BalanceController<T>::_UpdateLegCMD(ControlFSMData<T>& data) {
    LegControllerCommand<T>* cmd = data._legController->commands;
    auto& seResult = data._stateEstimator->getResult();
    for(int leg = 0; leg < 4; leg ++){
        if(data.userParameters->cmpc_gait == 4){
            Vec3<T> pDes,qDes;
            pDes <<  0, 0, -stand_height;
            computeLegAngle(*data._quadruped, pDes, &qDes);
            cmd[leg].qDes = qDes;
            cmd[leg].qdDes = Vec3<T>::Zero();
            cmd[leg].kpJoint << 100,0,0,0,50,0,0,0,50 ;
            cmd[leg].kdJoint << 0.2,0,0,0,0.2,0,0,0,0.2 ;       

        }
        else{
            cmd[leg].qDes = _pid_data.qDes[leg];
            cmd[leg].qdDes = _pid_data.qdDes[leg];
            cmd[leg].kpJoint = _pid_data.kpJoint[leg];
            cmd[leg].kdJoint = _pid_data.kdJoint[leg];
        }
    }
}

template class BalanceController<float>;
template class BalanceController<double>;
